/*
 * test_node.cpp
 *
 *  Created on: 16.09.2013
 *      Author: alex
 */

#include <octomap/octomap.h>

#include <bullet_rviz_test/test_node.h>
#include <stdlib.h>
#include <math.h>
#include <boost/make_shared.hpp>
#include <ros/ros.h>
#include <bullet/btBulletDynamicsCommon.h>
#include <bullet/btBulletCollisionCommon.h>
#include <bullet_rviz_test/rviz_motion_state.h>
#include <bullet_rviz_test/physics.h>
#include <bullet_rviz_test/physical_object.h>
#include <bullet_rviz_test/heightfield.h>
#include <bullet_rviz_test/heightfield_visualization.h>
#include <bullet_rviz_test/robot.h>
#include <bullet_rviz_test/robot_visualization.h>
//#include <traverse_analysis/traverse_analysis.h>

static const uint32_t GROUND_COLOR = 0x999988FFU;
static const size_t NUMBER_OF_ROBOTS_X = 5;
static const size_t NUMBER_OF_ROBOTS_Y = 5;
static const float WORST_DISTANCE = 0.2;
static const float WORST_ROTATION = 3.1415f / 4;

using namespace octomap;
using namespace std;

int main(int argc, char ** argv)
{
  return bullet_rviz_test::TestNode::main(argc, argv);
}

namespace bullet_rviz_test
{

TestNode::TestNode()
{
  reset_timer_ = nh_.createTimer(ros::Duration(5.0),
      &TestNode::resetTimerCallback, this);

  //traverse_analysis_ = boost::make_shared<traverse_analysis::TraverseAnalysis>(
  //    nh_);
  //traverse_analysis_->init();

  physics_ = boost::make_shared<Physics>();

  createMockHeightfield();
  //createHeightfield();

  physics_->addPhysicalObject(heightfield_);
  heightfield_visualization_ = boost::make_shared<HeightfieldVisualization>(
      heightfield_, GROUND_COLOR);

  for (size_t i = 0; i < NUMBER_OF_ROBOTS_X * NUMBER_OF_ROBOTS_Y; ++i)
  {
    RobotInfoPtr robot_info = boost::make_shared<RobotInfo>();

    robot_info->robot_ = boost::make_shared<Robot>();
    physics_->addPhysicalObject(robot_info->robot_);
    robot_info->robot_->registerMotionCallback(
        PhysicalObject::MotionCallback(
            boost::bind(&TestNode::robotMotionCallback, this,
                robot_info.get())));

    robot_info->visualization_ = boost::make_shared<RobotVisualization>(
        robot_info->robot_);

    robots_.push_back(robot_info);
  }

  resetRobots();
}

int TestNode::main(int argc, char ** argv)
{
  ros::init(argc, argv, "bullet_rviz_test");
  TestNode node;
  return node.run();
}

int TestNode::run()
{
  std::cout << __func__ << ": started" << std::endl;

  ros::Rate loop_rate(30);
  while (ros::ok())
  {
    ros::spinOnce();
    loop_rate.sleep();
    physics_->stepSimulation();
  }

  robots_.clear();
  heightfield_visualization_.reset();
  heightfield_.reset();

  return 0;
}

void TestNode::resetRobots()
{
  btVector3 heightfield_extents = heightfield_->getExtents();
  btVector3 heightfield_position =
      heightfield_->getRigidBody()->getWorldTransform().getOrigin();
  for (size_t y_index = 0; y_index < NUMBER_OF_ROBOTS_Y; ++y_index)
  {
    for (size_t x_index = 0; x_index < NUMBER_OF_ROBOTS_X; ++x_index)
    {
      size_t i = y_index * NUMBER_OF_ROBOTS_X + x_index;
      RobotPtr & robot = robots_[i]->robot_;

      btScalar x = heightfield_position.x()
          + heightfield_extents.x()
              * (-0.5f + (x_index + 0.5f) / NUMBER_OF_ROBOTS_X);
      btScalar y = heightfield_position.y()
          + heightfield_extents.y()
              * (-0.5f + (y_index + 0.5f) / NUMBER_OF_ROBOTS_Y);
      btScalar z = heightfield_->getHeightAtPosition(x, y);
      if (isnan(z) || z < (heightfield_position.z() - heightfield_extents.z()))
        z = heightfield_position.z() - heightfield_extents.z()
            - 2 * robot->getHalfExtents().z();
      else
        z += 2 * robot->getHalfExtents().z();

      z += 5.0;

      btVector3 position(x, y, z);
      robot->getRigidBody()->setWorldTransform(
          btTransform(btQuaternion(0, 0, 0, 1), position));
      robot->getRigidBody()->setLinearVelocity(btVector3(0, 0, 0));
      robot->getRigidBody()->setAngularVelocity(btVector3(0, 0, 0));
      robot->getRigidBody()->activate();

      robots_[i]->initial_position_ = btVector3(x, y,
          z - robot->getHalfExtents().z());
      //std::cout << "(" << x << ", " << y << ") ";
    }
    //std::cout << std::endl;
  }
}

void TestNode::computeCosts(RobotInfo * robot_info)
{
  btTransform world_transform =
      robot_info->robot_->getRigidBody()->getWorldTransform();
  btVector3 position = world_transform.getOrigin();

  btScalar moved_distance = position.distance(robot_info->initial_position_);

  float cost = moved_distance / WORST_DISTANCE
      + fabs(world_transform.getRotation().getAngle()) / WORST_ROTATION;

  robot_info->visualization_->setCost(cost);
}

void TestNode::resetTimerCallback(const ros::TimerEvent &)
{
  heightfield_visualization_->visualize();
  resetRobots();
}

void TestNode::createMockHeightfield()
{
//  btVector3 heightfield_extents(10, 9, 2);
//  size_t size_x = 30;
//  size_t size_y = 30;
//
//  float altitude_x = heightfield_extents.z() * 0.2f;
//  float angular_offset_x = -0.25 * 3.1415f;
//  float angular_velocity_x = 2.2 * 3.1415f / size_x;
//
//  float altitude_y = heightfield_extents.z() * 0.8f;
//  float angular_offset_y = 0.20 * 3.1415f;
//  float angular_velocity_y = 1.4 * 3.1415f / size_y;
//
//  float * data = new float[size_x * size_y];
//  for (size_t y_index = 0; y_index < size_y; ++y_index)
//  {
//    for (size_t x_index = 0; x_index < size_x; ++x_index)
//    {
//      data[y_index * size_x + x_index] = altitude_x
//          * sin(angular_offset_x + angular_velocity_x * x_index)
//          + altitude_y * sin(angular_offset_y + angular_velocity_y * y_index);
//    }
//  }
//
//  for (size_t y_index = size_y / 4; y_index < size_y * 3 / 4; ++y_index)
//  {
//    for (size_t x_index = size_x / 4; x_index < size_x * 3 / 4; ++x_index)
//    {
//      data[y_index * size_x + x_index] = -heightfield_extents.z() - 0.1;
//    }
//  }
//  heightfield_ = boost::make_shared<Heightfield>(heightfield_extents, data,
//      size_x, size_y);
//  delete[] data;

  OcTree tree("data/ssrr_arena_map_3.bt");

  double min_x, min_y, min_z;
  tree.getMetricMin(min_x, min_y, min_z);

  double max_x, max_y, max_z;
  tree.getMetricMax(max_x, max_y, max_z);

  btVector3 heightfield_extents(max_x-min_x, max_y-min_y, max_z-min_z);

  double node_size = tree.getResolution();
  std::cout << node_size << std::endl;

  std::cout << "grid x " << (max_x - min_x)/node_size << std::endl;
  std::cout << "grid y " << (max_y - min_y)/node_size << std::endl;

  std::cout << "room size: " << (max_x - min_x) << "/" << (max_y - min_y) << "/" << (max_z - min_z) << std::endl;

  int width = ceil((max_x - min_x)/node_size);
  int height = ceil((max_y - min_y)/node_size);

  // create height map
  float * data = new float[width * height];
  for (int i=0; i<width*height; i++)
    data[i] = min_z;

  for(OcTree::leaf_iterator it = tree.begin_leafs(),
         end=tree.end_leafs(); it!= end; ++it)
  {
    point3d center = it.getCoordinate();
    double current_size = it.getSize();

//    std::cout << it->getValue() << std::endl;

    double left_upper_x = center.x() - current_size/2.0 + node_size/2.0;
    double left_upper_y = center.y() - current_size/2.0 + node_size/2.0;

    for (double x = left_upper_x; x < left_upper_x + current_size; x += node_size)
    {
      for (double y = left_upper_y; y < left_upper_y + current_size; y += node_size)
      {
        int index_x = floor((x-min_x)/node_size);
        int index_y = floor((y-min_y)/node_size);

        double block_top = center.z() + node_size/2.0;

        // value > 0 means more than 50% (prob=exp(value)/(1+exp(value)))
        if (data[index_x + index_y*width] < block_top && it->getValue() > 0.0)
          data[index_x + index_y*width] = block_top;
      }
    }
  }


  heightfield_ = boost::make_shared<Heightfield>(heightfield_extents, data,
                                                 width, height);
  delete[] data;
}

//  for (size_t y_index = 0; y_index < size_y; ++y_index)
//  {
//    for (size_t x_index = 0; x_index < size_x; ++x_index)
//    {
//      data[y_index * size_x + x_index] = altitude_x
//          * sin(angular_offset_x + angular_velocity_x * x_index)
//      + altitude_y * sin(angular_offset_y + angular_velocity_y * y_index);
//    }
//  }

  // initialize elevation map
//  std::vector<std::vector<double> > elevation_map(width);
//  for (int i=0; i<width; i++)
//  {
//    elevation_map[i].resize(height);
//    for (int j=0; j<height; j++)
//    {
//      elevation_map[i][j] = -1000.0;
//    }
//  }
//
//  for(OcTree::leaf_iterator it = tree.begin_leafs(),
//         end=tree.end_leafs(); it!= end; ++it)
//  {
//    point3d center = it.getCoordinate();
//    double current_size = it.getSize();
//
//    int left_upper_x = center.x() - current_size/2.0 + node_size/2.0;
//    int left_upper_y = center.y() - current_size/2.0 + node_size/2.0;
//
//    for (double x = left_upper_x; x < left_upper_x + current_size; x += node_size)
//    {
//      for (double y = left_upper_y; y < left_upper_y + current_size; y += node_size)
//      {
//        int index_x = floor((x-min_x)/node_size);
//        int index_y = floor((y-min_y)/node_size);
//
//        if (elevation_map[index_x][index_y] < center.z())
//          elevation_map[index_x][index_y] = center.z();
//      }
//    }
//  }

/*
void TestNode::createHeightfield()
{
  traverse_analysis::HeightFieldPtr height_field =
      traverse_analysis_->getHeightfieldMap();

  btVector3 heightfield_extents(height_field->length, height_field->width, 3);

  float * data = new float[height_field->rows * height_field->cols];
  for (size_t y_index = 0; y_index < height_field->rows; ++y_index)
  {
    for (size_t x_index = 0; x_index < height_field->cols; ++x_index)
    {
      float value = height_field->height_map[y_index
          + height_field->rows * x_index];
      if (isnan(value))
        value = -heightfield_extents.z() - 0.1;
      data[y_index * height_field->cols + x_index] = value;
    }
  }

  heightfield_ = boost::make_shared<Heightfield>(heightfield_extents, data,
      height_field->cols, height_field->rows);
  delete[] data;
}
*/

void TestNode::robotMotionCallback(RobotInfo * robot_info)
{
  computeCosts(robot_info);
}

}
